#pragma once
#include <utility>

enum COLOR
{
	RED,
	BLACK
};

template<class K, class V>
struct RBTreeNode
{
    RBTreeNode(const pair<K, V>& val)
        :_left(nullptr)
        ,_right(nullptr)
        ,_parent(nullptr)
        , _kv(val)
        ,_color(RED)
    {}

	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;
	pair<K, V> _kv;
	COLOR _color;
};

template<class K, class V>
class RBTree
{
	typedef RBTreeNode<K, V> Node;

public:
    RBTree()
        :_root(nullptr) 
    {}

    bool Insert(const pair<K, V>& kv)
    {
        //第一次插入数据
        if (_root == nullptr)
        {
            //根节点只能是黑色
            _root = new Node(kv);
            _root->_color = BLACK;
            return true;
        }

        //找到合适的位置插入
        Node* parent = nullptr;
        Node* cur = _root;

        while (cur)
        {
            if (cur->_kv.first < kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                //key已经存在
                return false;
            }
        }

        //走到这里表示找到合适的位置插入
        cur = new Node(kv);
        cur->_color = RED;//新节点为红色

        if (parent->_kv.first < kv.first)
        {
            parent->_right = cur;
        }
        else
        {
            parent->_left = cur;
        }

        cur->_parent = parent;

        //红黑树的控制逻辑
        //父亲节点的颜色为黑色不需要调整
        //反之，需要调整
        while (parent && parent->_color == RED)
        {
            Node* grandFather = parent->_parent;
            if (parent == grandFather->_left)//考虑父亲在祖父的位置
            {
                Node* uncle = grandFather->_right;
                
                if (uncle && uncle->_color == RED)//uncle 存在且为红
                {
                    //将uncle和parent改成黑色，grandfather改成红色
                    uncle->_color = parent->_color = BLACK;
                    grandFather->_color = RED;

                    //如果祖父就是根节点，将祖父修改成黑色，退出即可
                    if (grandFather == _root)
                    {
                        grandFather->_color = BLACK;
                        break;
                    }
                    else //然后继续向上迭代处理
                    {
                        cur = grandFather;
                        parent = cur->_parent;
                    }
                }
                else //uncle不存在 或者 uncle存在且为黑
                {
                    if (parent->_left == cur) //单旋的场景
                    {
                        RotateR(grandFather);
                        parent->_color = BLACK;
                        grandFather->_color = RED;
                    }
                    else // LR双旋转
                    {
                        RotateL(parent);
                        RotateR(grandFather);
                        cur->_color = BLACK;
                        grandFather->_color = RED;
                    }

                    break; //旋转后直接退出即可
                }
            }
            else //parent == grandFather->_right
            {
                Node* uncle = grandFather->_left;

                if (uncle && uncle->_color == RED) //uncle 存在且为红
                {
                    uncle->_color = parent->_color = BLACK;
                    grandFather->_color = RED;

                    //如果祖父就是根节点，将祖父修改成黑色，退出即可
                    if (grandFather == _root)
                    {
                        grandFather->_color = BLACK;
                        break;
                    }
                    else //然后继续向上迭代处理
                    {
                        cur = grandFather;
                        parent = cur->_parent;
                    }
                }
                else //uncle不存在 或者 uncle存在且为黑
                {
                    if (parent->_right == cur) //单旋的场景
                    {
                        RotateL(grandFather);
                        parent->_color = BLACK;
                        grandFather->_color = RED;
                    }
                    else // RL双旋转
                    {
                        RotateR(parent);
                        RotateL(grandFather);
                        cur->_color = BLACK;
                        grandFather->_color = RED;
                    }

                    break; //旋转后直接退出即可
                }
            }
        }

        return true;
    }

    void RotateL(Node* parent)
    {
        //_rotateCount++;
        //将curLeft连接到parent的右子树中
        //将parent连接到cur的左子树中
        //修改cur、parent、curLeft的_parent

        Node* cur = parent->_right;
        Node* curLeft = cur->_left;
        parent->_right = curLeft;

        // curLeft可能为空
        if (curLeft)
            curLeft->_parent = parent;

        cur->_left = parent;

        //提前保存parent的_parent避免节点丢失
        Node* ppnode = parent->_parent;
        parent->_parent = cur;


        //边界情况，修改_root
        if (parent == _root)
        {
            _root = cur;
            cur->_parent = nullptr;
        }
        else
        {
            if (ppnode->_left == parent)
                ppnode->_left = cur;
            else
                ppnode->_right = cur;

            cur->_parent = ppnode;
        }
    }

    void RotateR(Node* parent)
    {
        //++_rotateCount;
        //将cur的右子树连接到parent的左子数
        //让parent连接到cur的右子树
        //修改parent、cur、curRight的_parent

        Node* cur = parent->_left;
        Node* curRight = cur->_right;

        parent->_left = curRight;

        //curRight可能为空
        if (curRight)
            curRight->_parent = parent;

        //避免parent->_parent节点丢失
        Node* ppnode = parent->_parent;

        parent->_parent = cur;
        cur->_right = parent;

        //边界情况
        if (parent == _root)
        {
            _root = cur;
            cur->_parent = nullptr;
        }
        else
        {
            if (ppnode->_left == parent)
                ppnode->_left = cur;
            else
                ppnode->_right = cur;

            cur->_parent = ppnode;
        }
    }

    bool checkcolor(Node* root, int blacknum, int benchmark)
    {
        //处理边界情况
        if (root == nullptr) 
        {
            //走到NIL节点，基准值与当前路径黑节点数量不一致
            if (blacknum != benchmark)
                return false;

            return true;
        }
        
        //统计当前路径的黑节点数量
        if (root->_color == BLACK)
            ++blacknum;

        //出现连续红节点的情况
        if (root->_color == RED && root->_parent && root->_parent->_color == RED)
        {
            cout << root->_kv.first << "-> 位置出现连续红节点" << endl;
            return false;
        }

        //子问题分治
        return checkcolor(root->_left, blacknum, benchmark) && checkcolor(root->_right, blacknum, benchmark);
    }

    bool IsBalance()
    {
        return IsBalance(_root);
    }
    
    bool IsBalance(Node* root)
    {
        //处理边界情况
        if (root == nullptr)
            return true;
        
        //根节点只能为黑色
        if (_root->_color == RED)
        {
            return false;
        }

        //以最左路径的黑节点的数量为基准值
        int benchmark = 0;
        Node* cur = root;
        while (cur)
        {
            if(cur->_color == BLACK)
                ++benchmark;
            cur = cur->_left;
        }

        return checkcolor(root, 0, benchmark);
    }

    int Height()
    {
        return Height(_root);
    }

    int Height(Node* root)
    {
        //处理边界情况
        if (root == nullptr)
            return 0;

        int leftH = Height(root->_left);
        int rightH = Height(root->_right);

        return leftH > rightH ? leftH + 1 : rightH + 1;
    }

private:
	Node* _root;

//public:
//    int _rotateCount = 0;
};